#!/usr/bin/env python3
#######################################################
#
#  Email   : 1450595701@qq.com
#
#  program describe:
#     This program is a tool, migrate OpenCloudOS8 to OpenCloudOS9.
#########################################################

import ast
import re
import os
import sys
import shutil
import argparse
import subprocess
from libs.common import ServiceCheck,RpmStatus
from libs.common import Requirements,BaseTools
from libs.common import InstRpmsInfo,InstRpmsName
from libs.common import ServiceStatus,RpmQaV,YumStatus
from libs.common import InstRpms,WriteData
from libs.variable import migrate_repostr
from libs.variable import services,config


# Set variables
log_dir = config["log_dir"]
yum_dir = config["yum_dir"]
data_dir = config["data_dir"]
work_dir = config["work_dir"]
kernel_pkg = ast.literal_eval(config["kernel_pkg"])
base_system = ast.literal_eval(config["base_system"])
conflicts_pkgs = ast.literal_eval(config["conflicts_pkgs"])
regex_solved_case = ast.literal_eval(config["regex_solved_case"])


def MigrateBeforeRepoDirSet():
    print("[INFO]: Delete the old repo and keep the third-party repo set to disabled.")
    if not os.path.isdir(yum_dir):
        print(f"[FAIL]: Yum dir '{yum_dir}' is not exist, please check!")
        sys.exit(1)
    subprocess.run('yum install -y yum-utils wget', shell=True)

    for f in os.listdir(f"{yum_dir}"):
        file_path = os.path.join(yum_dir + "/" + f)
        if os.path.isdir(file_path):
            continue
        if re.match('^OpenCloudOS.*\.repo$', f):
            os.remove(yum_dir + "/" + f)
            continue
        if re.fullmatch('switch-to-openclous-9.0.repo', f, re.IGNORECASE):
            continue
        with open(file_path, 'r') as extrepo:
            extrepo_content = extrepo.read()
            ''' Write install log to file '''
            WriteData(file_path+'.disabled', extrepo_content, mode="w")
            os.remove(file_path)


def EnableThirdPartyRepo():
    ''' Enable third-party repositories '''
    for ef in os.listdir(f"{yum_dir}"):
        efile_path = os.path.join(yum_dir + "/" + ef)
        if not os.path.isfile(efile_path):
            continue
        if re.match(f".*\.disabled$", ef):
            enabled_repo = ef.replace(".disabled", "")
            enabled_repo_path = os.path.join(yum_dir + "/" + enabled_repo)
            os.rename(efile_path, enabled_repo_path)


def InstO9DstRelasePkg():
    print("[INFO]: Remove Opencloud OS 8 dst RPMS and install Opencloud OS 9 dst RPMS")
    dst_release_list = ['opencloudos-repos', 'opencloudos-release']
    subprocess.run('rpm -e $(rpm -qa | grep \"^opencloudos\" | grep -E \"logos|release|repos|repo\") --nodeps', shell=True)
    print("migrate_repostr:", migrate_repostr)
    with open(f"{yum_dir}/switch-to-openclous-9.0.repo", 'w') as migrate_repostr_data:
        migrate_repostr_data.write(migrate_repostr)
    try:
        subprocess.check_output("yumdownloader "+' '.join(dst_release_list), shell=True)
    except Exception as e:
        print(":Error]: Failed to download the Opencloud OS 9 release Packages.")
        sys.exit(1)
    subprocess.run('rpm -e $(rpm -q gpg-pubkey --qf "%{NAME}-%{VERSION}-%{RELEASE} %{PACKAGER}\\n" | grep OpenCloudOS | awk \'{print $1}\')', shell=True)
    subprocess.run(f"rpm -ivh opencloudos-* --force --nodeps", shell=True)
    os.remove(f"{yum_dir}/switch-to-openclous-9.0.repo")


def GetDiskInfo(diststr):
    ''' Get system disk info '''
    dev_name = ""
    part_name = ""
    length = len(string)
    for c in range(length-1, -1, -1):
        if not string[c].isdigit():
            if string.find('nvme') != -1:
                dev_name = string[0:c]
                part_num = string[c+1:length]
            else:
                dev_name = string[0:c+1]
                part_num = string[c+1:length]
                break
    return dev_name,part_num

def SetEfibootmgr():
    print("[INFO]: Current system is uefi, add boot option to boot manager.")
    disk_name = ""
    subprocess.run('which efibootmgr > /dev/null 2>&1 || yum install -y efibootmgr', shell=True)
    disk_name = subprocess.check_output('mount | grep /boot/efi | awk \'{print $1}\'', shell=True)
    disk_name = str(disk_name, 'utf-8')
    disk_name = disk_name.split('\n')[0]
    dev_name,part_num = GetDiskInfo(disk_name)
    if dev_name == "" or part_num == "":
        print("[Fail]: Parse /boot/efi disk info failed, update boot loader failed.")
        sys.exit(1)
    cmd = ""
    arch = platform.machine()
    if arch == "x86_64":
        cmd = 'efibootmgr -c -d ' + dev_name + ' -p ' + part_num + ' -l "/EFI/opencloudos/shimaa64.efi" -L "Migrate To OpenCloudOS 9"'
    elif arch == "aarch64":
        cmd = 'efibootmgr -c -d ' + dev_name + ' -p ' + part_num + ' -l "/EFI/opencloudos/shimx64.efi" -L "Migrate To OpenCloudOS 9"'
    efi_rst = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    if efi_rst.returncode != 0:
        print("[Fail]: Use efibootmgr update boot loader failed.")
        sys.exit(1)

def UpdateGrubCfg():
    ''' Update grub config '''
    if os.path.isdir('/sys/firmware/efi'):
        subprocess.run('grub2-mkconfig -o /boot/efi/EFI/opencloudos/grub.cfg', shell=True)
        SetEfibootmgr()
    else:
        subprocess.run('grub2-mkconfig -o /boot/grub2/grub.cfg', shell=True)
        try:
            subprocess.check_call('test -L /boot/grub2/grubenv', shell=True)
            subprocess.run('mv /boot/grub2/grubenv /boot/grub2/grubenv-bak', shell=True)
            subprocess.run('cat /boot/grub2/grubenv-bak > /boot/grub2/grubenv', shell=True)
        except:
            pass

def MigrateToOc9():
    print("[INFO]: Begin to migtate Opencloud OS 8 to Opencloud OS 9")
    subprocess.run("systemctl stop firewalld", shell=True)
    for conflicts_pkgs_case in conflicts_pkgs:
        print(f"[INFO]: Remove the conflicts rpm with {conflicts_pkgs_case}")
        subprocess.run(f"rpm -e {conflicts_pkgs_case} --nodeps", shell=True)
    InstRpms(*base_system, opts="--allowerasing",prompt="y")

    ''' Solve ssh service not  allow root user to login by passwd '''
    subprocess.run("cp -arf /etc/ssh/sshd_config /etc/ssh/sshd_config.old", shell=True)
    InstRpms("openssh-server")
    subprocess.run(['sudo', 'sed', '-i', 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/g', '/etc/ssh/sshd_config'])
    subprocess.run(['sudo', 'systemctl', 'restart', 'sshd'])

    InstRpms("distro-sync", opts="", prompt="y")
    ''' Problems "Regex version", it will be solved below. '''
    InstRpms(*regex_solved_case)


    ''' Enable thethird-party repos to install the missing pkgs '''
    if args.sync_all_pkg:
        with open(f"{data_dir}/migrate_before_rpmsname.txt", 'r') as before_names_str:
            before_names_set = set(before_names_str_list.strip() for before_names_str_list in before_names_str)
        ''' Get the latest missing pkgs '''
        after_rpmnsame_set = InstRpmsName("after")
        latest_miss_rpms_set = before_names_set.difference(after_rpmnsame_set)
        print("latest_miss_rpms_set:",latest_miss_rpms_set)
        print("[INFO]: Enable the third-party mirrors to install the missing pkgs!")
        EnableThirdPartyRepo()

        if not YumStatus():
            print("[Warn]: May be the third-party mirrors were already tried, skip to install the missing pkgs!")
            return
        else:
            if latest_miss_rpms_set:
                latest_miss_rpms_code = InstRpms(*latest_miss_rpms_set, opts="--exclude=kernel-*", prompt="y")
                if latest_miss_rpms_code != 0:
                    InstRpms(*latest_miss_rpms_set, opts="--exclude=kernel-* --skip-broken", prompt="y")

    ''' Makr sure kernel pkg is exist! '''
    print("[INFO]: Install kernel pkgs of " + ' '.join(kernel_pkg))
    for install_kernel in kernel_pkg:
        if not RpmStatus(install_kernel):
            kernel_pkg_code = InstRpms(install_kernel ,prompt="y")
            if kernel_pkg_code != 0:
                print("[FAIL]: Failed to install about kernel Package!")
                sys.exit(1)

    
    ''' Updating initrd '''
    if os.access('/usr/libexec/plymouth/plymouth-update-initrd', os.X_OK):
        subprocess.run('/usr/libexec/plymouth/plymouth-update-initrd')
    ''' Update grub config '''
    UpdateGrubCfg()


def OpenCloudOS8ToOpenCloudOS9():
    ''' The main run funtion to migrate OpenCloudOS 8 to OpenCloudOS 9 steps '''
    if not args.continue_run:
        Requirements()
        BaseTools()
        InstRpmsInfo("before")
        InstRpmsName("before")
        ServiceStatus("before")
        RpmQaV("before")
        MigrateBeforeRepoDirSet()
        InstO9DstRelasePkg()
    MigrateToOc9()
    InstRpmsInfo("after")
    InstRpmsName("after")
    ServiceStatus("after")
    RpmQaV("after")
    ServiceCheck()
    LatestMsg()

def InstallUserModeRpms(*install_params):
    print("install_params",install_params)
    print(f"Executing hanshu02 function with argument: {install_params}")
    InstRpms(*install_params, opts="--exclude=kernel-*", prompt="y")

def LatestMsg():
    '''Out put the lastet info'''
    latest_msg = f"""\
[Guidance prompt]
----------------------------------------
1. Check the changes of installation package files before and after migration.
        cd /var/tmp/migrate/data/
        diff -Nurp migrate_before_rpmV.txt migrate_after_rpmV.txt

2. Check the changes of installation package before and after migration.
        cd /var/tmp/migrate/data/
        diff -Nurp migrate_before_rpmsname.txt migrate_after_rpmsname.txt

3. Check the changes of system services before and after migration.
        cd /var/tmp/migrate/data/
        diff -Nurp migrate_before_service_names.txt migrate_after_service_names.txt

4. Check the changes of complete information of system installation package before and after migration.
        cd /var/tmp/migrate/data/
        cat migrate_before_rpmsinfo.json;  cat migrate_after_rpmsinfo.json
---------------------------------------
Recommended Action:
Migrate complete. OpenCloudOS 9.0 recommends rebooting this system.
"""
    print(latest_msg)

def main():
    global args
    os.environ["LC_ALL"] = "en_US.UTF-8"
    parser = argparse.ArgumentParser()
    parser = argparse.ArgumentParser(description='Process some arguments.')
    parser.add_argument('-i', '--install', type=str, nargs='?', const="", help='Install the RPMS Packages log to /tmp/tmp/migrate/logs')
    parser.add_argument('-a', '--all', action='store_const', const="y", dest='sync_all_pkg', help="Sync all packages!")
    parser.add_argument('-c', '--continue', action='store_const', const="y", dest='continue_run', help="Resolve the conflicts then continue!")
    args = parser.parse_args()

    # Run tools:
    if not args.install:
        OpenCloudOS8ToOpenCloudOS9()
    elif args.install and not args.install.strip():
        OpenCloudOS8ToOpenCloudOS9()
    elif args.install and args.install.strip():
        InstallUserModeRpms(args.install)

if __name__ == "__main__":
    main()
